{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using custom data configuration codeparrot--codeparrot-clean-valid-826c6fd8b27e5523\n",
      "Found cached dataset json (/Users/loubnabenallal/.cache/huggingface/datasets/codeparrot___json/codeparrot--codeparrot-clean-valid-826c6fd8b27e5523/0.0.0/e6070c77f18f01a5ad4551a8b7edfba20b8438b7cad4d94e6ad9378022ce4aab)\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline\n",
    "\n",
    "ds = load_dataset(\"codeparrot/codeparrot-clean-valid\", split=\"train\")\n",
    "ds = ds.select([i for i in range(100)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Person name detection using an NER model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will test [bert-base-NER](https://huggingface.co/dslim/bert-base-NER) and [CodeBert finetuned](https://huggingface.co/mrm8488/codebert-base-finetuned-stackoverflow-ner) on an NER dataset on Stackoverflow to detect Person Names in a subset of our Python dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### BERT-Base-NER"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(\"dslim/bert-base-NER\")\n",
    "model = AutoModelForTokenClassification.from_pretrained(\"dslim/bert-base-NER\")\n",
    "\n",
    "nlp_bert = pipeline(\"ner\", model=model, tokenizer=tokenizer, aggregation_strategy=\"simple\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 149,
   "metadata": {},
   "outputs": [],
   "source": [
    "example = \"\\n# @python3.9\\n#@ejonath\\n# This work is contributed by Marc Dupont and  Elizabeth Johnson  and is part of the Satellite project\\n# Python program for implementation of Bubble Sort\\n\\ndef bubbleSort(arr):\\n    n = len(arr)\\n    # Traverse through all array elements\\n    # Credits to Taylor Jonahthan for this section\\n    for i in range(n):\\n\\n        # Last i elements are already in place\\n        for j in range(0, n-i-1):\\n\\n            # traverse the array from 0 to n-i-1\\n            # Swap if the element found is greater\\n            # than the next element\\n            if arr[j] > arr[j+1]:\\n                arr[j], arr[j+1] = arr[j+1], arr[j]\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 150,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "# @python3.9\n",
      "#@ejonath\n",
      "# This work is contributed by Marc Dupont and  Elizabeth Johnson  and is part of the Satellite project\n",
      "# Python program for implementation of Bubble Sort\n",
      "\n",
      "def bubbleSort(arr):\n",
      "    n = len(arr)\n",
      "    # Traverse through all array elements\n",
      "    # Credits to Taylor Jonahthan for this section\n",
      "    for i in range(n):\n",
      "\n",
      "        # Last i elements are already in place\n",
      "        for j in range(0, n-i-1):\n",
      "\n",
      "            # traverse the array from 0 to n-i-1\n",
      "            # Swap if the element found is greater\n",
      "            # than the next element\n",
      "            if arr[j] > arr[j+1]:\n",
      "                arr[j], arr[j+1] = arr[j+1], arr[j]\n"
     ]
    }
   ],
   "source": [
    "print(example)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 151,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'entity_group': 'PER', 'score': 0.9996812, 'word': 'Marc Dupont', 'start': 54, 'end': 65}\n",
      "{'entity_group': 'PER', 'score': 0.9995078, 'word': 'Elizabeth Johnson', 'start': 71, 'end': 88}\n",
      "{'entity_group': 'PER', 'score': 0.9963441, 'word': 'Taylor Jonahthan', 'start': 276, 'end': 292}\n"
     ]
    }
   ],
   "source": [
    "ner_results = nlp_bert(example)\n",
    "\n",
    "for e in ner_results:\n",
    "    if e[\"entity_group\"] ==\"PER\":\n",
    "        print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### CodeBert StackOverflow NER"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer_so = AutoTokenizer.from_pretrained(\"mrm8488/codebert-base-finetuned-stackoverflow-ner\")\n",
    "model_so = AutoModelForTokenClassification.from_pretrained(\"mrm8488/codebert-base-finetuned-stackoverflow-ner\")\n",
    "nlp_so = pipeline(\"ner\", model=model_so, tokenizer=tokenizer_so, aggregation_strategy=\"simple\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 152,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'entity_group': 'User_Name', 'score': 0.9836406, 'word': 'e', 'start': 16, 'end': 17}\n",
      "{'entity_group': 'User_Name', 'score': 0.6776985, 'word': 'jonath', 'start': 17, 'end': 23}\n",
      "{'entity_group': 'User_Name', 'score': 0.98596036, 'word': ' Marc', 'start': 54, 'end': 58}\n",
      "{'entity_group': 'User_Name', 'score': 0.8043251, 'word': ' Elizabeth', 'start': 71, 'end': 80}\n",
      "{'entity_group': 'User_Name', 'score': 0.969354, 'word': ' Taylor', 'start': 276, 'end': 282}\n"
     ]
    }
   ],
   "source": [
    "ner_results_so = nlp_so(example)\n",
    "\n",
    "for e in ner_results_so:\n",
    "    if e[\"entity_group\"] == \"User_Name\":\n",
    "        print(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'entity': 'B-User_Name', 'score': 0.98323464, 'index': 15, 'word': 'ĠMarc', 'start': 44, 'end': 48}\n",
      "{'entity': 'I-User_Name', 'score': 0.75678843, 'index': 16, 'word': 'ĠDup', 'start': 49, 'end': 52}\n",
      "{'entity': 'B-User_Name', 'score': 0.6773796, 'index': 20, 'word': 'ĠElizabeth', 'start': 61, 'end': 70}\n",
      "{'entity': 'B-User_Name', 'score': 0.9776358, 'index': 75, 'word': 'ĠTaylor', 'start': 266, 'end': 272}\n",
      "{'entity': 'I-User_Name', 'score': 0.42680034, 'index': 76, 'word': 'ĠJonah', 'start': 273, 'end': 278}\n"
     ]
    }
   ],
   "source": [
    "# without pipeline aggregation\n",
    "nlp_so_no_agg = pipeline(\"ner\", model=model_so, tokenizer=tokenizer_so)\n",
    "res = nlp_so_no_agg(example)\n",
    "\n",
    "for e in res:\n",
    "    if \"User_Name\" in e[\"entity\"]:\n",
    "        print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Apply on some Python files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "metadata": {},
   "outputs": [],
   "source": [
    "def detect_names_bert(example):\n",
    "    entities = nlp_bert(example[\"content\"])\n",
    "    names = [sample for sample in entities if sample[\"entity_group\"] == \"PER\"]\n",
    "    return {\"names\": names}\n",
    "\n",
    "def detect_names_so(example):\n",
    "    entities = nlp_so(example[\"content\"])\n",
    "    names = [sample for sample in entities if sample[\"entity_group\"] == \"User_Name\"]\n",
    "    return {\"names\": names}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Detection with BERT-Base-NER"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "755c11fb0f984af4b9e406d58be618b2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/100 [00:00<?, ?ex/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ds_bert = ds.map(detect_names_bert)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Names detected in 21 elements which are [1, 8, 14, 17, 23, 34, 37, 39, 45, 46, 49, 55, 69, 70, 72, 76, 88, 89, 92, 93, 98]\n"
     ]
    }
   ],
   "source": [
    "output = [i for i in range(len(ds_bert)) if ds_bert[i][\"names\"]]\n",
    "print(f\"Names detected in {len(output)} elements which are {output}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sample 1:\n",
      " [{'end': 197, 'entity_group': 'PER', 'score': 0.9986525774002075, 'start': 195, 'word': 'Mi'}, {'end': 199, 'entity_group': 'PER', 'score': 0.8829349875450134, 'start': 197, 'word': '##ki'}, {'end': 209, 'entity_group': 'PER', 'score': 0.9712410569190979, 'start': 199, 'word': '##o Hirabaya'}]\n",
      "Names are ['Mi', '##ki', '##o Hirabaya']\n",
      "\n",
      "Sample 8:\n",
      " [{'end': 31, 'entity_group': 'PER', 'score': 0.9820157885551453, 'start': 23, 'word': 'Luc Saff'}]\n",
      "Names are ['Luc Saff']\n",
      "\n",
      "Sample 14:\n",
      " [{'end': 359, 'entity_group': 'PER', 'score': 0.9671221971511841, 'start': 358, 'word': 'C'}, {'end': 467, 'entity_group': 'PER', 'score': 0.875209629535675, 'start': 466, 'word': 'C'}, {'end': 849, 'entity_group': 'PER', 'score': 0.8578362464904785, 'start': 848, 'word': 'C'}]\n",
      "Names are ['C', 'C', 'C']\n",
      "\n",
      "Sample 17:\n",
      " [{'end': 28, 'entity_group': 'PER', 'score': 0.9939444661140442, 'start': 26, 'word': 'Ce'}, {'end': 42, 'entity_group': 'PER', 'score': 0.9917498230934143, 'start': 28, 'word': '##dric Bellegard'}]\n",
      "Names are ['Ce', '##dric Bellegard']\n",
      "\n",
      "Sample 23:\n",
      " [{'end': 921, 'entity_group': 'PER', 'score': 0.4104381203651428, 'start': 920, 'word': 'O'}, {'end': 1061, 'entity_group': 'PER', 'score': 0.6051611304283142, 'start': 1060, 'word': 'O'}, {'end': 1135, 'entity_group': 'PER', 'score': 0.8839671611785889, 'start': 1120, 'word': 'Mario Steinhoff'}]\n",
      "Names are ['O', 'O', 'Mario Steinhoff']\n",
      "\n",
      "Sample 34:\n",
      " [{'end': 699, 'entity_group': 'PER', 'score': 0.982388973236084, 'start': 685, 'word': 'Andrew Bentley'}]\n",
      "Names are ['Andrew Bentley']\n",
      "\n",
      "Sample 37:\n",
      " [{'end': 130, 'entity_group': 'PER', 'score': 0.9921714067459106, 'start': 119, 'word': 'Jason Jones'}]\n",
      "Names are ['Jason Jones']\n",
      "\n",
      "Sample 39:\n",
      " [{'end': 377, 'entity_group': 'PER', 'score': 0.710799515247345, 'start': 369, 'word': 'Walter P'}, {'end': 382, 'entity_group': 'PER', 'score': 0.7480440139770508, 'start': 379, 'word': '##car'}]\n",
      "Names are ['Walter P', '##car']\n",
      "\n",
      "Sample 45:\n",
      " [{'end': 236, 'entity_group': 'PER', 'score': 0.9990164041519165, 'start': 225, 'word': 'Marc Abramo'}, {'end': 429, 'entity_group': 'PER', 'score': 0.9991106390953064, 'start': 412, 'word': 'Benjamin Peterson'}]\n",
      "Names are ['Marc Abramo', 'Benjamin Peterson']\n",
      "\n",
      "Sample 46:\n",
      " [{'end': 59, 'entity_group': 'PER', 'score': 0.9345136880874634, 'start': 49, 'word': 'Luc Saffre'}]\n",
      "Names are ['Luc Saffre']\n",
      "\n",
      "Sample 49:\n",
      " [{'end': 52, 'entity_group': 'PER', 'score': 0.9813946485519409, 'start': 38, 'word': 'Travis Painter'}]\n",
      "Names are ['Travis Painter']\n",
      "\n",
      "Sample 55:\n",
      " [{'end': 232, 'entity_group': 'PER', 'score': 0.9086423516273499, 'start': 221, 'word': 'Victor Olay'}, {'end': 946, 'entity_group': 'PER', 'score': 0.9522541761398315, 'start': 938, 'word': 'Victor O'}]\n",
      "Names are ['Victor Olay', 'Victor O']\n",
      "\n",
      "Sample 69:\n",
      " [{'end': 430, 'entity_group': 'PER', 'score': 0.9983365535736084, 'start': 429, 'word': 'R'}, {'end': 433, 'entity_group': 'PER', 'score': 0.9716659784317017, 'start': 432, 'word': 'J'}, {'end': 441, 'entity_group': 'PER', 'score': 0.9301518797874451, 'start': 435, 'word': 'Gowers'}, {'end': 444, 'entity_group': 'PER', 'score': 0.9993035793304443, 'start': 443, 'word': 'M'}, {'end': 451, 'entity_group': 'PER', 'score': 0.9185934066772461, 'start': 444, 'word': '. Linke'}, {'end': 454, 'entity_group': 'PER', 'score': 0.9993711113929749, 'start': 453, 'word': 'J'}, {'end': 463, 'entity_group': 'PER', 'score': 0.9445327520370483, 'start': 454, 'word': '. Barnoud'}, {'end': 466, 'entity_group': 'PER', 'score': 0.998609185218811, 'start': 465, 'word': 'T'}, {'end': 479, 'entity_group': 'PER', 'score': 0.8763769865036011, 'start': 468, 'word': 'J. E. Reddy'}, {'end': 482, 'entity_group': 'PER', 'score': 0.9990938901901245, 'start': 481, 'word': 'M'}, {'end': 491, 'entity_group': 'PER', 'score': 0.8568034172058105, 'start': 482, 'word': '. N. Melo'}, {'end': 494, 'entity_group': 'PER', 'score': 0.9992411136627197, 'start': 493, 'word': 'S'}, {'end': 497, 'entity_group': 'PER', 'score': 0.8912100791931152, 'start': 494, 'word': '. L'}, {'end': 505, 'entity_group': 'PER', 'score': 0.998450756072998, 'start': 499, 'word': 'Seyler'}, {'end': 510, 'entity_group': 'PER', 'score': 0.9989950060844421, 'start': 509, 'word': 'D'}, {'end': 513, 'entity_group': 'PER', 'score': 0.7343029975891113, 'start': 510, 'word': '. L'}, {'end': 521, 'entity_group': 'PER', 'score': 0.9904451966285706, 'start': 515, 'word': 'Dotson'}, {'end': 524, 'entity_group': 'PER', 'score': 0.9994157552719116, 'start': 523, 'word': 'J'}, {'end': 534, 'entity_group': 'PER', 'score': 0.8871639370918274, 'start': 524, 'word': '. Domanski'}, {'end': 537, 'entity_group': 'PER', 'score': 0.999458372592926, 'start': 536, 'word': 'S'}, {'end': 546, 'entity_group': 'PER', 'score': 0.8947599530220032, 'start': 537, 'word': '. Buchoux'}, {'end': 549, 'entity_group': 'PER', 'score': 0.9982117414474487, 'start': 548, 'word': 'I'}, {'end': 552, 'entity_group': 'PER', 'score': 0.8505293130874634, 'start': 551, 'word': 'M'}, {'end': 560, 'entity_group': 'PER', 'score': 0.9938455820083618, 'start': 554, 'word': 'Kenney'}, {'end': 567, 'entity_group': 'PER', 'score': 0.9995276927947998, 'start': 566, 'word': 'O'}, {'end': 578, 'entity_group': 'PER', 'score': 0.9711530804634094, 'start': 569, 'word': 'Beckstein'}, {'end': 675, 'entity_group': 'PER', 'score': 0.9977560639381409, 'start': 674, 'word': 'S'}, {'end': 685, 'entity_group': 'PER', 'score': 0.8529825210571289, 'start': 677, 'word': 'Benthall'}, {'end': 691, 'entity_group': 'PER', 'score': 0.9985602498054504, 'start': 690, 'word': 'S'}, {'end': 700, 'entity_group': 'PER', 'score': 0.8084614872932434, 'start': 693, 'word': 'Rostrup'}, {'end': 811, 'entity_group': 'PER', 'score': 0.998878538608551, 'start': 810, 'word': 'N'}, {'end': 820, 'entity_group': 'PER', 'score': 0.9980390071868896, 'start': 813, 'word': 'Michaud'}, {'end': 828, 'entity_group': 'PER', 'score': 0.9882857203483582, 'start': 821, 'word': 'Agrawal'}, {'end': 831, 'entity_group': 'PER', 'score': 0.998974084854126, 'start': 830, 'word': 'E'}, {'end': 843, 'entity_group': 'PER', 'score': 0.8114681243896484, 'start': 831, 'word': '. J. Denning'}, {'end': 846, 'entity_group': 'PER', 'score': 0.9986012578010559, 'start': 845, 'word': 'T'}, {'end': 856, 'entity_group': 'PER', 'score': 0.8595272898674011, 'start': 846, 'word': '. B. Woolf'}, {'end': 863, 'entity_group': 'PER', 'score': 0.9994939565658569, 'start': 862, 'word': 'O'}, {'end': 874, 'entity_group': 'PER', 'score': 0.959717869758606, 'start': 865, 'word': 'Beckstein'}, {'end': 955, 'entity_group': 'PER', 'score': 0.40819162130355835, 'start': 954, 'word': 'J'}]\n",
      "Names are ['R', 'J', 'Gowers', 'M', '. Linke', 'J', '. Barnoud', 'T', 'J. E. Reddy', 'M', '. N. Melo', 'S', '. L', 'Seyler', 'D', '. L', 'Dotson', 'J', '. Domanski', 'S', '. Buchoux', 'I', 'M', 'Kenney', 'O', 'Beckstein', 'S', 'Benthall', 'S', 'Rostrup', 'N', 'Michaud', 'Agrawal', 'E', '. J. Denning', 'T', '. B. Woolf', 'O', 'Beckstein', 'J']\n",
      "\n",
      "Sample 70:\n",
      " [{'end': 16, 'entity_group': 'PER', 'score': 0.9975435137748718, 'start': 15, 'word': 'R'}, {'end': 18, 'entity_group': 'PER', 'score': 0.9954125285148621, 'start': 16, 'word': '##ox'}, {'end': 23, 'entity_group': 'PER', 'score': 0.7910621762275696, 'start': 18, 'word': '##ane P'}]\n",
      "Names are ['R', '##ox', '##ane P']\n",
      "\n",
      "Sample 72:\n",
      " [{'end': 215, 'entity_group': 'PER', 'score': 0.9861843585968018, 'start': 200, 'word': 'Alexandre Defen'}]\n",
      "Names are ['Alexandre Defen']\n",
      "\n",
      "Sample 76:\n",
      " [{'end': 139, 'entity_group': 'PER', 'score': 0.9568077921867371, 'start': 127, 'word': 'Nicolas Wack'}]\n",
      "Names are ['Nicolas Wack']\n",
      "\n",
      "Sample 88:\n",
      " [{'end': 74, 'entity_group': 'PER', 'score': 0.9972798824310303, 'start': 68, 'word': 'Adrian'}, {'end': 88, 'entity_group': 'PER', 'score': 0.7989827394485474, 'start': 74, 'word': '##o Monteiro Mar'}]\n",
      "Names are ['Adrian', '##o Monteiro Mar']\n",
      "\n",
      "Sample 89:\n",
      " [{'end': 140, 'entity_group': 'PER', 'score': 0.650191605091095, 'start': 138, 'word': 'Po'}, {'end': 152, 'entity_group': 'PER', 'score': 0.8787997961044312, 'start': 150, 'word': 'Po'}, {'end': 397, 'entity_group': 'PER', 'score': 0.5624960660934448, 'start': 395, 'word': 'Po'}, {'end': 714, 'entity_group': 'PER', 'score': 0.9567288160324097, 'start': 712, 'word': 'Po'}]\n",
      "Names are ['Po', 'Po', 'Po', 'Po']\n",
      "\n",
      "Sample 92:\n",
      " [{'end': 202, 'entity_group': 'PER', 'score': 0.944946825504303, 'start': 190, 'word': 'Frank Cavigg'}, {'end': 233, 'entity_group': 'PER', 'score': 0.9643455743789673, 'start': 221, 'word': 'Frank Cavigg'}]\n",
      "Names are ['Frank Cavigg', 'Frank Cavigg']\n",
      "\n",
      "Sample 93:\n",
      " [{'end': 84, 'entity_group': 'PER', 'score': 0.9994305968284607, 'start': 82, 'word': 'Ko'}, {'end': 93, 'entity_group': 'PER', 'score': 0.9674811363220215, 'start': 84, 'word': '##vid Goyal'}]\n",
      "Names are ['Ko', '##vid Goyal']\n",
      "\n",
      "Sample 98:\n",
      " [{'end': 56, 'entity_group': 'PER', 'score': 0.9969924092292786, 'start': 44, 'word': 'Richard Paus'}]\n",
      "Names are ['Richard Paus']\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for i, sample in enumerate(ds_bert):\n",
    "    if sample[\"names\"]:\n",
    "        print(f\"Sample {i}:\\n {sample['names']}\")\n",
    "        print(f\"Names are {[e['word'] for e in sample['names'] if e]}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Detection with CodeBERT-Stackoverflow-NER"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2ed40e90d498443f9da2f3b40717c221",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/100 [00:00<?, ?ex/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ds_so = ds.map(detect_names_so)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Names detected in 26 elements which are [1, 8, 13, 15, 17, 21, 23, 26, 34, 37, 38, 45, 46, 49, 52, 55, 70, 72, 76, 84, 88, 92, 93, 94, 96, 98]\n"
     ]
    }
   ],
   "source": [
    "output = [i for i in range(len(ds_so)) if ds_so[i][\"names\"]]\n",
    "print(f\"Names detected in {len(output)} elements which are {output}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 123,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sample 1:\n",
      " [{'end': 200, 'entity_group': 'User_Name', 'score': 0.6459169983863831, 'start': 195, 'word': ' Mikio'}]\n",
      "Names are [' Mikio']\n",
      "\n",
      "Sample 8:\n",
      " [{'end': 31, 'entity_group': 'User_Name', 'score': 0.7196319699287415, 'start': 23, 'word': ' Luc Saff'}]\n",
      "Names are [' Luc Saff']\n",
      "\n",
      "Sample 13:\n",
      " [{'end': 65, 'entity_group': 'User_Name', 'score': 0.718497633934021, 'start': 61, 'word': ' Kids'}]\n",
      "Names are [' Kids']\n",
      "\n",
      "Sample 15:\n",
      " [{'end': 732, 'entity_group': 'User_Name', 'score': 0.8268138766288757, 'start': 731, 'word': ' d'}, {'end': 738, 'entity_group': 'User_Name', 'score': 0.22002415359020233, 'start': 732, 'word': 'ulwich'}]\n",
      "Names are [' d', 'ulwich']\n",
      "\n",
      "Sample 17:\n",
      " [{'end': 29, 'entity_group': 'User_Name', 'score': 0.892275333404541, 'start': 26, 'word': ' Ced'}, {'end': 48, 'entity_group': 'User_Name', 'score': 0.7896466255187988, 'start': 45, 'word': 'ced'}, {'end': 51, 'entity_group': 'User_Name', 'score': 0.43889954686164856, 'start': 48, 'word': 'ric'}]\n",
      "Names are [' Ced', 'ced', 'ric']\n",
      "\n",
      "Sample 21:\n",
      " [{'end': 139, 'entity_group': 'User_Name', 'score': 0.7978324294090271, 'start': 139, 'word': ' '}, {'end': 146, 'entity_group': 'User_Name', 'score': 0.9026286602020264, 'start': 140, 'word': ' NigelB'}]\n",
      "Names are [' ', ' NigelB']\n",
      "\n",
      "Sample 23:\n",
      " [{'end': 61, 'entity_group': 'User_Name', 'score': 0.7056369781494141, 'start': 60, 'word': ' f'}, {'end': 1135, 'entity_group': 'User_Name', 'score': 0.7441920638084412, 'start': 1120, 'word': ' Mario Steinhoff'}]\n",
      "Names are [' f', ' Mario Steinhoff']\n",
      "\n",
      "Sample 26:\n",
      " [{'end': 719, 'entity_group': 'User_Name', 'score': 0.9294928312301636, 'start': 717, 'word': ' iv'}, {'end': 722, 'entity_group': 'User_Name', 'score': 0.5376254916191101, 'start': 719, 'word': 'ana'}, {'end': 728, 'entity_group': 'User_Name', 'score': 0.5403547883033752, 'start': 724, 'word': 'ihal'}]\n",
      "Names are [' iv', 'ana', 'ihal']\n",
      "\n",
      "Sample 34:\n",
      " [{'end': 691, 'entity_group': 'User_Name', 'score': 0.5979154109954834, 'start': 685, 'word': 'Andrew'}, {'end': 726, 'entity_group': 'User_Name', 'score': 0.26092109084129333, 'start': 723, 'word': 'rew'}]\n",
      "Names are ['Andrew', 'rew']\n",
      "\n",
      "Sample 37:\n",
      " [{'end': 130, 'entity_group': 'User_Name', 'score': 0.7241939902305603, 'start': 119, 'word': ' Jason Jones'}, {'end': 133, 'entity_group': 'User_Name', 'score': 0.5918320417404175, 'start': 132, 'word': 'j'}, {'end': 137, 'entity_group': 'User_Name', 'score': 0.23216299712657928, 'start': 133, 'word': 'ason'}, {'end': 142, 'entity_group': 'User_Name', 'score': 0.4017036557197571, 'start': 141, 'word': 'j'}, {'end': 156, 'entity_group': 'User_Name', 'score': 0.24113701283931732, 'start': 154, 'word': 'es'}]\n",
      "Names are [' Jason Jones', 'j', 'ason', 'j', 'es']\n",
      "\n",
      "Sample 38:\n",
      " [{'end': 284, 'entity_group': 'User_Name', 'score': 0.7965254783630371, 'start': 283, 'word': ' D'}]\n",
      "Names are [' D']\n",
      "\n",
      "Sample 45:\n",
      " [{'end': 174, 'entity_group': 'User_Name', 'score': 0.72637939453125, 'start': 166, 'word': 'msabramo'}, {'end': 240, 'entity_group': 'User_Name', 'score': 0.8863577842712402, 'start': 225, 'word': ' Marc Abramowitz'}, {'end': 252, 'entity_group': 'User_Name', 'score': 0.9774307608604431, 'start': 249, 'word': ' mar'}, {'end': 253, 'entity_group': 'User_Name', 'score': 0.458333283662796, 'start': 252, 'word': 'c'}, {'end': 257, 'entity_group': 'User_Name', 'score': 0.8579809069633484, 'start': 254, 'word': 'mar'}, {'end': 429, 'entity_group': 'User_Name', 'score': 0.9176520109176636, 'start': 412, 'word': ' Benjamin Peterson'}]\n",
      "Names are ['msabramo', ' Marc Abramowitz', ' mar', 'c', 'mar', ' Benjamin Peterson']\n",
      "\n",
      "Sample 46:\n",
      " [{'end': 59, 'entity_group': 'User_Name', 'score': 0.6349635124206543, 'start': 49, 'word': ' Luc Saffre'}]\n",
      "Names are [' Luc Saffre']\n",
      "\n",
      "Sample 49:\n",
      " [{'end': 52, 'entity_group': 'User_Name', 'score': 0.6837303638458252, 'start': 38, 'word': ' Travis Painter'}]\n",
      "Names are [' Travis Painter']\n",
      "\n",
      "Sample 52:\n",
      " [{'end': 57, 'entity_group': 'User_Name', 'score': 0.7014908194541931, 'start': 53, 'word': ' Meng'}, {'end': 65, 'entity_group': 'User_Name', 'score': 0.41586849093437195, 'start': 63, 'word': 'gu'}, {'end': 75, 'entity_group': 'User_Name', 'score': 0.5847292542457581, 'start': 68, 'word': 'mxgnene'}]\n",
      "Names are [' Meng', 'gu', 'mxgnene']\n",
      "\n",
      "Sample 55:\n",
      " [{'end': 230, 'entity_group': 'User_Name', 'score': 0.9039267897605896, 'start': 221, 'word': ' Victor Ol'}, {'end': 268, 'entity_group': 'User_Name', 'score': 0.7904767990112305, 'start': 261, 'word': ' volayaf'}]\n",
      "Names are [' Victor Ol', ' volayaf']\n",
      "\n",
      "Sample 70:\n",
      " [{'end': 21, 'entity_group': 'User_Name', 'score': 0.7377451658248901, 'start': 15, 'word': ' Roxane'}]\n",
      "Names are [' Roxane']\n",
      "\n",
      "Sample 72:\n",
      " [{'end': 207, 'entity_group': 'User_Name', 'score': 0.9391152262687683, 'start': 200, 'word': ' Alexand'}]\n",
      "Names are [' Alexand']\n",
      "\n",
      "Sample 76:\n",
      " [{'end': 136, 'entity_group': 'User_Name', 'score': 0.7292793989181519, 'start': 127, 'word': ' Nicolas W'}, {'end': 142, 'entity_group': 'User_Name', 'score': 0.6803852319717407, 'start': 141, 'word': 'w'}]\n",
      "Names are [' Nicolas W', 'w']\n",
      "\n",
      "Sample 84:\n",
      " [{'end': 55, 'entity_group': 'User_Name', 'score': 0.8204792737960815, 'start': 53, 'word': ' ng'}, {'end': 57, 'entity_group': 'User_Name', 'score': 0.47535043954849243, 'start': 55, 'word': 'ak'}, {'end': 59, 'entity_group': 'User_Name', 'score': 0.33769160509109497, 'start': 57, 'word': 'an'}]\n",
      "Names are [' ng', 'ak', 'an']\n",
      "\n",
      "Sample 88:\n",
      " [{'end': 74, 'entity_group': 'User_Name', 'score': 0.8504272103309631, 'start': 68, 'word': ' Adrian'}, {'end': 109, 'entity_group': 'User_Name', 'score': 0.9683395028114319, 'start': 105, 'word': ' Amit'}, {'end': 119, 'entity_group': 'User_Name', 'score': 0.7104912400245667, 'start': 115, 'word': 'amix'}]\n",
      "Names are [' Adrian', ' Amit', 'amix']\n",
      "\n",
      "Sample 92:\n",
      " [{'end': 195, 'entity_group': 'User_Name', 'score': 0.9707217812538147, 'start': 190, 'word': ' Frank'}, {'end': 226, 'entity_group': 'User_Name', 'score': 0.9730982184410095, 'start': 221, 'word': ' Frank'}, {'end': 241, 'entity_group': 'User_Name', 'score': 0.7681715488433838, 'start': 237, 'word': 'fcav'}]\n",
      "Names are [' Frank', ' Frank', 'fcav']\n",
      "\n",
      "Sample 93:\n",
      " [{'end': 85, 'entity_group': 'User_Name', 'score': 0.7950971722602844, 'start': 82, 'word': ' Kov'}, {'end': 98, 'entity_group': 'User_Name', 'score': 0.8132345676422119, 'start': 95, 'word': 'kov'}]\n",
      "Names are [' Kov', 'kov']\n",
      "\n",
      "Sample 94:\n",
      " [{'end': 53, 'entity_group': 'User_Name', 'score': 0.9459285736083984, 'start': 52, 'word': ' B'}, {'end': 60, 'entity_group': 'User_Name', 'score': 0.5108811259269714, 'start': 53, 'word': 'itergia'}, {'end': 732, 'entity_group': 'User_Name', 'score': 0.6695467233657837, 'start': 731, 'word': 's'}, {'end': 763, 'entity_group': 'User_Name', 'score': 0.9529682397842407, 'start': 759, 'word': ' Quan'}, {'end': 772, 'entity_group': 'User_Name', 'score': 0.6479182839393616, 'start': 770, 'word': 'qu'}]\n",
      "Names are [' B', 'itergia', 's', ' Quan', 'qu']\n",
      "\n",
      "Sample 96:\n",
      " [{'end': 48, 'entity_group': 'User_Name', 'score': 0.8862028121948242, 'start': 41, 'word': ' Niphlod'}, {'end': 57, 'entity_group': 'User_Name', 'score': 0.7316068410873413, 'start': 50, 'word': 'niphlod'}, {'end': 980, 'entity_group': 'User_Name', 'score': 0.368248850107193, 'start': 975, 'word': 'Niphl'}, {'end': 988, 'entity_group': 'User_Name', 'score': 0.3562377095222473, 'start': 984, 'word': 'niph'}]\n",
      "Names are [' Niphlod', 'niphlod', 'Niphl', 'niph']\n",
      "\n",
      "Sample 98:\n",
      " [{'end': 51, 'entity_group': 'User_Name', 'score': 0.9701380133628845, 'start': 44, 'word': ' Richard'}]\n",
      "Names are [' Richard']\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for i, sample in enumerate(ds_so):\n",
    "    if sample[\"names\"]:\n",
    "        print(f\"Sample {i}:\\n {sample['names']}\")\n",
    "        print(f\"Names are {[e['word'] for e in sample['names'] if e]}\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 124,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'repo_name': Value(dtype='string', id=None),\n",
       " 'path': Value(dtype='string', id=None),\n",
       " 'copies': Value(dtype='string', id=None),\n",
       " 'size': Value(dtype='string', id=None),\n",
       " 'content': Value(dtype='string', id=None),\n",
       " 'license': Value(dtype='string', id=None),\n",
       " 'hash': Value(dtype='int64', id=None),\n",
       " 'line_mean': Value(dtype='float64', id=None),\n",
       " 'line_max': Value(dtype='int64', id=None),\n",
       " 'alpha_frac': Value(dtype='float64', id=None),\n",
       " 'autogenerated': Value(dtype='bool', id=None),\n",
       " 'names': [{'end': Value(dtype='int64', id=None),\n",
       "   'entity_group': Value(dtype='string', id=None),\n",
       "   'score': Value(dtype='float32', id=None),\n",
       "   'start': Value(dtype='int64', id=None),\n",
       "   'word': Value(dtype='string', id=None)}]}"
      ]
     },
     "execution_count": 124,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds_so.features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 156,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0e4be3c007334ac3a9ca9325607ce9e3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/100 [00:00<?, ?ex/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def name_column(example):\n",
    "    person_names = []\n",
    "    if example[\"names\"]:\n",
    "        person_names = [e[\"word\"] for e in example[\"names\"] if e]\n",
    "    return {\"person_names\": person_names}\n",
    "\n",
    "ds_so_new = ds_so.map(name_column, remove_columns=[\"repo_name\", \"copies\", \"hash\", 'line_mean', 'line_max', 'alpha_frac', 'autogenerated'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 159,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bc1ac0d9349c4b40a54605da1657169f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Pushing dataset shards to the dataset hub:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ds_so_new.push_to_hub(\"names_detection_codebert\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 160,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "47f8b799e2904f638a336e9ba66f5627",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/100 [00:00<?, ?ex/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ds_bert_new = ds_bert.map(name_column, remove_columns=[\"repo_name\", \"copies\", \"hash\", 'line_mean', 'line_max', 'alpha_frac', 'autogenerated'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 161,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3e20c8bbbe3848c7afdb87a81be5e50d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Pushing dataset shards to the dataset hub:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ds_bert_new.push_to_hub(\"names_detection_bert_ner\")"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "fd8fde6f83dada9276d12fdb71d773558994168ed1b3bea457b8db38c02aa2e1"
  },
  "kernelspec": {
   "display_name": "Python 3.10.4 ('venv')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
